# -*- coding: utf-8 -*-

"""
图像拼板工具
将实际的矩形图像拼接到指定大小的板材上，用于打印
宽度限制：1米(1000mm)
间距：10mm
支持图像旋转
"""

import os
import json
import argparse
import math
from typing import List, Dict, Tuple
from PIL import Image, ImageDraw, ImageFont
import numpy as np

class ImageItem:
    """图像物品类，用于表示要拼板的图像"""
    def __init__(self, image_path, index=None):
        self.image_path = image_path
        self.index = index
        self.x = 0  # 放置位置的x坐标
        self.y = 0  # 放置位置的y坐标
        self.rotated = False  # 是否旋转
        
        # 打开图像获取尺寸
        with Image.open(image_path) as img:
            self.width, self.height = img.size
            self.area = self.width * self.height
    
    def rotate(self):
        """旋转图像，交换宽高"""
        self.width, self.height = self.height, self.width
        self.rotated = not self.rotated
    
    def get_dimensions(self):
        """获取图像的尺寸"""
        return self.width, self.height
    
    def __repr__(self):
        return f"ImageItem(path={self.image_path}, width={self.width}, height={self.height}, index={self.index}, rotated={self.rotated})"

class ImageNesting:
    """图像拼板类，实现图像装箱算法"""
    def __init__(self, width_limit=1000, height_limit=None, spacing=10, dpi=300):
        self.width_limit = width_limit  # 宽度限制，默认1000mm
        self.height_limit = height_limit  # 高度限制，默认无限制
        self.spacing = spacing  # 间距，默认10mm
        self.dpi = dpi  # 分辨率，默认300DPI
        self.mm_to_px = dpi / 25.4  # 毫米到像素的转换因子
        self.px_to_mm = 25.4 / dpi  # 像素到毫米的转换因子
        self.items = []  # 所有要拼板的图像
        self.nesting_plan = []  # 拼板计划
        self.total_width_mm = 0  # 总宽度(mm)
        self.total_height_mm = 0  # 总高度(mm)
    
    def load_images_from_directory(self, directory_path):
        """从目录加载图像文件"""
        try:
            # 获取目录中的所有图像文件
            image_extensions = ['.png', '.jpg', '.jpeg', '.bmp', '.tiff']
            image_files = []
            
            # 需要忽略的文件名
            ignored_files = ['page_1_original_mask.png', 'page_1_full_page.png', 'page_1_annotated_mask.png']
            
            # 遍历目录及其子目录查找图像文件
            for root, dirs, files in os.walk(directory_path):
                for file in files:
                    # 检查文件是否是图像且不是需要忽略的文件
                    if (any(file.lower().endswith(ext) for ext in image_extensions) and 
                        file not in ignored_files):
                        image_files.append(os.path.join(root, file))
            
            # 按文件名排序
            image_files.sort()
            
            # 创建图像物品
            for idx, image_path in enumerate(image_files):
                try:
                    item = ImageItem(image_path, index=idx)
                    self.items.append(item)
                except Exception as e:
                    print(f"无法加载图像 {image_path}: {str(e)}")
            
            print(f"成功加载 {len(self.items)} 个图像")
            return True
        except Exception as e:
            print(f"加载图像时出错: {str(e)}")
            return False
    
    def load_images_from_json(self, json_path):
        """从JSON文件加载图像路径"""
        try:
            with open(json_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                
            # 从JSON中提取图像路径
            image_paths = []
            if 'pages' in data:
                for page in data['pages']:
                    if 'detected_rectangles' in page:
                        for idx, rect in enumerate(page['detected_rectangles']):
                            # 构造图像路径（假设图像保存在processed_images_organized目录下）
                            pdf_name = os.path.basename(json_path).split('_rectangle_results')[0]
                            image_path = os.path.join('processed_images_organized', pdf_name, 
                                                    f'page_{page.get("page_number", 1)}_rectangle_{idx+1}.png')
                            if os.path.exists(image_path):
                                image_paths.append(image_path)
                            else:
                                print(f"警告：图像文件不存在: {image_path}")
            
            # 创建图像物品
            for idx, image_path in enumerate(image_paths):
                try:
                    item = ImageItem(image_path, index=idx)
                    self.items.append(item)
                except Exception as e:
                    print(f"无法加载图像 {image_path}: {str(e)}")
            
            print(f"成功加载 {len(self.items)} 个图像")
            return True
        except Exception as e:
            print(f"加载JSON文件时出错: {str(e)}")
            return False
    
    def sort_items(self):
        """对图像进行排序，大的图像先放置"""
        # 按面积降序排序
        self.items.sort(key=lambda x: x.area, reverse=True)
    
    def optimize_rotation(self):
        """优化图像旋转，使宽度尽可能适合宽度限制"""
        width_limit_px = self.width_limit * self.mm_to_px  # 转换为像素
        
        for item in self.items:
            # 如果图像宽度超过限制，旋转它
            if item.width > width_limit_px and item.height <= width_limit_px:
                item.rotate()
            # 如果两种方向都超过限制，选择宽度较小的方向
            elif item.width > width_limit_px and item.height > width_limit_px:
                if item.width > item.height:
                    item.rotate()
    
    def nest_images(self):
        """执行图像拼板算法"""
        if not self.items:
            print("没有可拼板的图像")
            return False
        
        # 排序图像
        self.sort_items()
        
        # 优化旋转
        self.optimize_rotation()
        
        # 转换为像素单位
        width_limit_px = self.width_limit * self.mm_to_px
        spacing_px = self.spacing * self.mm_to_px
        
        # 初始化拼板区域
        shelves = []  # 货架列表
        current_y_px = 0
        
        # 对于每个图像，尝试放置到现有的货架，或者创建新货架
        for item in self.items:
            placed = False
            
            # 尝试放置到现有货架
            for shelf in shelves:
                # 尝试两种方向
                for _ in range(2):
                    # 检查是否可以放置
                    required_width = item.width
                    if len(shelf['items']) > 0:
                        required_width += spacing_px
                    
                    if required_width <= shelf['remaining_width']:
                        # 计算放置位置
                        if len(shelf['items']) > 0:
                            item.x = shelf['items'][-1]['x'] + shelf['items'][-1]['width'] + spacing_px
                        else:
                            item.x = 0
                        item.y = shelf['y']
                        
                        # 记录到货架
                        shelf['items'].append({
                            'x': item.x,
                            'y': item.y,
                            'width': item.width,
                            'height': item.height,
                            'rotated': item.rotated,
                            'image_path': item.image_path,
                            'index': item.index
                        })
                        
                        # 更新剩余宽度
                        if len(shelf['items']) > 1:
                            shelf['remaining_width'] -= spacing_px
                        shelf['remaining_width'] -= item.width
                        
                        placed = True
                        break
                    # 如果不能放置，尝试旋转
                    item.rotate()
                
                if placed:
                    break
            
            # 如果没有找到合适的货架，创建新货架
            if not placed:
                # 确保图像宽度不超过限制
                if item.width > width_limit_px:
                    # 尝试旋转
                    item.rotate()
                    if item.width > width_limit_px:
                        print(f"警告：图像 {item.index} 无法放置，尺寸过大")
                        continue
                
                # 创建新货架，高度为图像高度加上间距
                new_shelf = {
                    'y': current_y_px,
                    'height': item.height,
                    'remaining_width': width_limit_px,
                    'items': []
                }
                
                # 放置图像到新货架
                item.x = 0
                item.y = current_y_px
                
                new_shelf['items'].append({
                    'x': item.x,
                    'y': item.y,
                    'width': item.width,
                    'height': item.height,
                    'rotated': item.rotated,
                    'image_path': item.image_path,
                    'index': item.index
                })
                
                new_shelf['remaining_width'] -= item.width
                shelves.append(new_shelf)
                
                # 更新当前Y位置
                current_y_px += item.height + spacing_px
        
        # 计算总尺寸（减去最后一个间距）
        if shelves:
            self.total_height_mm = (current_y_px - spacing_px) * self.px_to_mm
        self.total_width_mm = self.width_limit
        
        # 保存拼板计划
        self.nesting_plan = shelves
        
        print(f"拼板完成，总尺寸: {self.total_width_mm:.2f}mm x {self.total_height_mm:.2f}mm")
        
        # 输出每个图像的位置信息
        for shelf in shelves:
            for item in shelf['items']:
                print(f"图像 {item['index']}: 位置=({item['x']*self.px_to_mm:.2f}mm, {item['y']*self.px_to_mm:.2f}mm), \
                      尺寸=({item['width']*self.px_to_mm:.2f}mm x {item['height']*self.px_to_mm:.2f}mm), 旋转={item['rotated']}")
        
        return True
    
    def create_nesting_image(self, output_path, draw_grid=True, draw_labels=True):
        """创建拼板图像"""
        if not self.nesting_plan:
            print("没有拼板计划可创建图像")
            return False
        
        try:
            # 转换为像素单位
            width_limit_px = self.width_limit * self.mm_to_px
            total_height_px = int(self.total_height_mm * self.mm_to_px)
            
            # 创建白色背景图像
            nesting_image = Image.new('RGB', (int(width_limit_px), total_height_px), color='white')
            draw = ImageDraw.Draw(nesting_image)
            
            # 尝试加载字体
            try:
                font = ImageFont.truetype("/System/Library/Fonts/PingFang.ttc", 20)
            except:
                font = ImageFont.load_default()
                print("警告：无法加载中文字体，使用默认字体")
            
            # 绘制每个图像
            for shelf in self.nesting_plan:
                for item_info in shelf['items']:
                    try:
                        # 打开原始图像
                        with Image.open(item_info['image_path']) as img:
                            # 如果需要旋转，旋转图像
                            if item_info['rotated']:
                                img = img.rotate(90, expand=True)
                            
                            # 确保图像尺寸匹配
                            if img.width != int(item_info['width']) or img.height != int(item_info['height']):
                                img = img.resize((int(item_info['width']), int(item_info['height'])), Image.LANCZOS)
                            
                            # 将图像粘贴到拼板上
                            nesting_image.paste(img, (int(item_info['x']), int(item_info['y'])))
                            
                            # 绘制边框
                            draw.rectangle([
                                (int(item_info['x']), int(item_info['y'])), 
                                (int(item_info['x'] + item_info['width']), int(item_info['y'] + item_info['height']))
                            ], outline='black', width=2)
                            
                            # 绘制标签
                            if draw_labels:
                                label = f"{item_info['index']}"
                                draw.text(
                                    (int(item_info['x'] + 10), int(item_info['y'] + 10)), 
                                    label, 
                                    fill='red', 
                                    font=font
                                )
                    except Exception as e:
                        print(f"处理图像 {item_info['image_path']} 时出错: {str(e)}")
            
            # 标注相邻矩形之间的间距
            for shelf in self.nesting_plan:
                if len(shelf['items']) > 1:
                    for i in range(len(shelf['items']) - 1):
                        item1 = shelf['items'][i]
                        item2 = shelf['items'][i + 1]
                        
                        # 计算间距（像素）
                        gap_px = item2['x'] - (item1['x'] + item1['width'])
                        gap_mm = gap_px * self.px_to_mm
                        
                        # 计算标注位置
                        x_pos = int(item1['x'] + item1['width'] + gap_px / 2)
                        y_pos = int((item1['y'] + item2['y']) / 2)
                        
                        # 绘制间距标注线
                        draw.line([(x_pos, y_pos - 20), (x_pos, y_pos + 20)], fill='blue', width=1)
                        
                        # 绘制间距数值
                        gap_text = f"{gap_mm:.1f}mm"
                        
                        # 使用更现代的方法计算文本尺寸
                        try:
                            # 对于较新版本的PIL/Pillow
                            bbox = draw.textbbox((0, 0), gap_text, font=font)
                            text_width = bbox[2] - bbox[0]
                            text_height = bbox[3] - bbox[1]
                        except AttributeError:
                            # 尝试使用textlength作为备选
                            text_width = draw.textlength(gap_text, font=font)
                            # 估计高度，可能不够准确，但作为后备方案
                            text_height = font.size if hasattr(font, 'size') else 20
                        
                        draw.rectangle([
                            (x_pos - text_width/2 - 5, y_pos - text_height/2 - 5),
                            (x_pos + text_width/2 + 5, y_pos + text_height/2 + 5)
                        ], fill='white', outline='blue', width=1)
                        draw.text(
                            (x_pos - text_width/2, y_pos - text_height/2),
                            gap_text,
                            fill='blue',
                            font=font
                        )
            
            # 绘制网格
            if draw_grid:
                # 绘制垂直线
                for x in range(0, int(width_limit_px), int(50 * self.mm_to_px)):  # 每50mm一条线
                    draw.line([(x, 0), (x, total_height_px)], fill='lightgray', width=1)
                # 绘制水平线
                for y in range(0, total_height_px, int(50 * self.mm_to_px)):  # 每50mm一条线
                    draw.line([(0, y), (width_limit_px, y)], fill='lightgray', width=1)
            
            # 添加标题
            title = f"图像拼板结果 - 尺寸: {self.total_width_mm:.2f}mm x {self.total_height_mm:.2f}mm ({self.width_limit}mm限制)"
            draw.text((20, 10), title, fill='black', font=font)
            
            # 保存图像
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            nesting_image.save(output_path)
            print(f"拼板图像已保存到: {output_path}")
            
            return output_path
        except Exception as e:
            print(f"创建拼板图像时出错: {str(e)}")
            return False
    
    def export_nesting_plan(self, output_path):
        """导出拼板计划到JSON文件"""
        if not self.nesting_plan:
            print("没有拼板计划可导出")
            return False
        
        try:
            # 准备导出数据
            export_data = {
                "width_limit_mm": self.width_limit,
                "total_width_mm": self.total_width_mm,
                "total_height_mm": self.total_height_mm,
                "spacing_mm": self.spacing,
                "dpi": self.dpi,
                "images": []
            }
            
            # 添加每个图像的信息
            for shelf in self.nesting_plan:
                for item_info in shelf['items']:
                    export_data["images"].append({
                        "index": item_info['index'],
                        "image_path": item_info['image_path'],
                        "x_mm": item_info['x'] * self.px_to_mm,
                        "y_mm": item_info['y'] * self.px_to_mm,
                        "width_mm": item_info['width'] * self.px_to_mm,
                        "height_mm": item_info['height'] * self.px_to_mm,
                        "rotated": item_info['rotated']
                    })
            
            # 保存到文件
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(export_data, f, ensure_ascii=False, indent=2)
            
            print(f"拼板计划已导出到: {output_path}")
            return True
        except Exception as e:
            print(f"导出拼板计划时出错: {str(e)}")
            return False


def main():
    """主函数，提供命令行接口"""
    parser = argparse.ArgumentParser(description='图像拼板工具 - 将实际图像拼接到指定大小的板材上')
    
    # 输入选项（互斥组）
    input_group = parser.add_mutually_exclusive_group(required=True)
    input_group.add_argument('--directory-path', type=str, help='图像所在目录路径')
    input_group.add_argument('--json-path', type=str, help='矩形结果JSON文件路径')
    
    # 输出选项
    parser.add_argument('--output-image', type=str, required=True, help='拼板结果图像输出路径')
    parser.add_argument('--export-plan', type=str, help='拼板计划JSON导出路径')
    
    # 拼板参数
    parser.add_argument('--width-limit', type=float, default=1000, help='拼板宽度限制(mm)，默认值为1000mm')
    parser.add_argument('--spacing', type=float, default=10, help='图像间距(mm)，默认值为10mm')
    parser.add_argument('--dpi', type=int, default=300, help='输出图像分辨率，默认值为300DPI')
    parser.add_argument('--no-grid', action='store_true', help='不绘制网格')
    parser.add_argument('--no-labels', action='store_true', help='不绘制图像标签')
    
    args = parser.parse_args()
    
    # 创建拼板实例
    nesting = ImageNesting(
        width_limit=args.width_limit,
        spacing=args.spacing,
        dpi=args.dpi
    )
    
    # 加载图像数据
    success = False
    if args.directory_path:
        success = nesting.load_images_from_directory(args.directory_path)
    elif args.json_path:
        success = nesting.load_images_from_json(args.json_path)
    
    if not success:
        print("加载图像失败，程序退出")
        return 1
    
    # 执行拼板
    if not nesting.nest_images():
        print("拼板失败，程序退出")
        return 1
    
    # 创建拼板图像
    if not nesting.create_nesting_image(
        args.output_image,
        draw_grid=not args.no_grid,
        draw_labels=not args.no_labels
    ):
        print("创建拼板图像失败，程序退出")
        return 1
    
    # 导出拼板计划
    if args.export_plan:
        nesting.export_nesting_plan(args.export_plan)
    
    return 0

if __name__ == '__main__':
    exit_code = main()
    import sys
    sys.exit(exit_code)